
#include <cxcore.h>
#include "median.hh"

#include "utils_shape.hh"


typedef enum  RobustFunc {HUBER, GM,CAUCHY,TUKEY,GAUSS,EXPO}RobustFunc ;

const double TUKEY_c = 4.6851;
const double HUBER_c = 1.345;
const double CAUCHY_c = 2.3849;
const double EXPO_c = 2.9846;
const double GM_c = 0.6745;


CvMat *projection_over_ST(const CvMat *shape, PDModel& pdm)
{

    const CvMat *s0 = pdm.mean;

    CvMat *proj = cvCreateMat(4,1, shape->type);
    cvZero(proj);
    CvMat *sh = cvCloneMat(shape);

    //get the centre of the shape
    
    CvScalar val;
    CvMat submat, row;

    for(int ai=0; ai<2; ai++){
	cvGetRow(sh, &submat, ai);
	val = cvAvg(&submat);
	cvSubS(&submat, val, &submat);
	cvmSet(proj, ai+2, 0, val.val[0]);
    }

    //project over the mean shape
    double q1 = cvDotProduct(sh, s0);
    
    CvMat *s_ = cvCreateMat(s0->rows, s0->cols, shape->type);
    cvGetRow(s0, &row, 0);
    cvGetRow(s_, &submat, 1);
    cvCopy(&row, &submat);

    cvGetRow(s0, &row, 1);
    cvGetRow(s_, &submat, 0);
    cvCopy(&row, &submat);
    cvScale(&submat, &submat, -1.);
    
    double q2 = cvDotProduct(sh, s_);

    cvReleaseMat(&s_);
    cvReleaseMat(&sh);
    
    cvmSet(proj,0, 0, q1-1);
    cvmSet(proj,1, 0, q2);

    return proj;

}

CvMat *projection_over_shape_base(const CvMat *sh, PDModel& pdm)
{

    const CvMat *s0 = pdm.mean;
    const CvMat *shVects = pdm.vectors;


    CvMat *proj = cvCreateMat(1, shVects->rows,sh->type);
    cvZero(proj);
    CvMat sh1d, s01d;

    cvReshape(sh, &sh1d, 0, 1);
    cvReshape(s0, &s01d, 0, 1);

    cvProjectPCA(&sh1d, &s01d, shVects, proj);

    cvTranspose(proj, proj);

    return proj;
}

void remove_ST(CvMat *sh, CvMat *st_p)
{

    CvMat submat;

    for(int ai=0; ai<2; ai++){
	cvGetRow(sh, &submat, ai);
	cvSubS(&submat, cvScalar(cvmGet(st_p,ai+2,0)), &submat);
    }

    CvMat *M = cvCreateMat(2,2,sh->type);

    cvmSet(M, 0,0, 1+cvmGet(st_p, 0,0));
    cvmSet(M, 1,1, 1+cvmGet(st_p, 0,0));
    cvmSet(M, 0,1, -1*cvmGet(st_p, 1,0));
    cvmSet(M, 1,0, cvmGet(st_p, 1,0));

    cvInvert(M, M);

    cvMatMul(M, sh, sh);
    
    cvReleaseMat(&M);

}

CvMat * instantiate_shape(PDModel& pdm, 
			  const CvMat *p_params, const CvMat *st_params) throw (std::invalid_argument)
{
    const CvMat *s0 = pdm.mean;
    const CvMat *shVects = pdm.vectors;

    if(CV_IS_MAT(p_params) && (p_params->cols*p_params->rows)!= shVects->rows)
	throw std::invalid_argument("The number of parameters and vectors is not the same");

    CvMat *sh = cvCloneMat(s0);
    CvMat sh1D;

    cvReshape(sh, &sh1D, 0, 1);

    if(CV_IS_MAT(p_params)){
	for(int ai = 0; ai< shVects->rows; ai++){
	    CvMat vec;
	    cvGetRow(shVects, &vec, ai);
	    cvAddWeighted(&sh1D, 1, &vec, cvGetReal1D(p_params, ai), 0, &sh1D);
	}
    }

    if(st_params){
	CvMat *homog, submat;	
	CvMat *mat = cvCreateMat(2,3, s0->type);

	cvmSet(mat,0,0,cvGetReal1D(st_params,0)+1);
	cvmSet(mat,0,1,-1*cvGetReal1D(st_params,1));
	cvmSet(mat,0,2,cvGetReal1D(st_params,2));
	cvmSet(mat,1,0,cvGetReal1D(st_params,1));
	cvmSet(mat,1,1,cvGetReal1D(st_params,0)+1);
	cvmSet(mat,1,2,cvGetReal1D(st_params,3));
		
	homog = cvCreateMat(3, sh->cols, s0->type);
	cvSet(homog, cvRealScalar(1.));
	cvGetRows(homog, &submat, 0, 2, 1);
	cvCopy(sh, &submat);

	cvMatMul(mat, homog, sh);

	cvReleaseMat(&mat);
	cvReleaseMat(&homog);
    }

    return sh;
}

//returns a 1xN matrix
CvMat *get_residues(const CvMat *shape, PDModel& pdm, 
		    const CvMat *p_param, const CvMat* st_params)
{

    CvMat *sh = instantiate_shape(pdm, p_param, st_params);
    CvMat *aux = cvCreateMat(1, 2*sh->cols, sh->type);
    CvMat mat;

    cvSub(sh, shape, sh);

    cvReshape(sh, &mat, 0, 1);
    cvCopy(&mat, aux);

    cvReleaseMat(&sh);

    return aux;
}

CvMat *get_robust_weight(const CvMat *residues, RobustFunc func = HUBER)
{

    CvMat *weight;
    CvMat *err_abs;
    double med, scale;
    
    //get the median
    err_abs= cvCreateMat(residues->rows, residues->cols, residues->type);    
    cvAbs(residues, err_abs);
    med = median(err_abs);

    weight = cvCreateMat(residues->rows, residues->cols, residues->type);

    //using GM
    if (func == GM){
	scale = med * GM_c;
	
	cvPow(err_abs, weight, 2);
	cvAddS(weight,cvScalar(scale*scale),weight);
	cvPow(weight, weight,2);
	//check if the any element of weight == 0
	CvMat *masc = NULL;
	int k = cvCountNonZero(weight) - 
	    weight->width*weight->height;

	if (k != 0){
	    masc = cvCreateMat(residues->rows, residues->cols, CV_8U);
	    cvCmpS(weight, 0., masc, CV_CMP_LE);
	    cvSet(weight, cvScalar(1),masc);
	}	    

	cvDiv(NULL, weight,weight, scale*scale*scale*scale);

	if(masc){
	    cvSet(weight,cvScalar(1),masc);
	    cvReleaseMat(&masc);
	}
		
    }
    //using Huber function
    else{
	scale = med * HUBER_c;
	cvDiv(NULL,err_abs, weight,scale);
	
	CvMat *masc = cvCreateMat(residues->rows, residues->cols, CV_8U);
	cvCmpS (err_abs, scale,masc,CV_CMP_LE);
	cvSet(weight,cvScalar(1), masc);
	    
	cvReleaseMat(&masc);
	
    }

    cvReleaseMat(&err_abs);
    
    return weight;

}

CvMat *calculate_robust_hessian(const CvMat *residue, const CvMat *vectors)
{

    CvMat *hessian = cvCreateMat(vectors->rows, vectors->rows, vectors->type);
    CvMat *temp1 = cvCreateMat(1, vectors->rows, vectors->type);

    cvZero(hessian);

    for(int ai=0; ai < vectors->cols; ai++){
	CvMat submat;
	cvGetCol(vectors, &submat, ai);
	cvTranspose(&submat, temp1);	

	cvGEMM(&submat, temp1, cvGetReal1D(residue, ai), hessian, 1, hessian);
    }


    cvReleaseMat(&temp1);
    return hessian;

}


CvMat *robust_projection_over_ST(const CvMat *shape, PDModel& pdm, 
				 const CvMat *params, const CvMat *st_p0)
{


    CvMat *proj = cvCloneMat(st_p0);

    CvMat submat, col, sh1d;
    cvReshape(shape, &sh1d, 0, 1);
    int nparams = st_p0->cols * st_p0->rows; //should be 4
   

    CvMat *ihess = cvCreateMat(nparams,nparams, shape->type);
    CvMat *term = cvCreateMat(nparams,1,shape->type);
    CvMat *update = cvCreateMat(nparams,1,shape->type);

    //build the base->PDM
    
    CvMat *base = cvCreateMat(nparams, pdm.vectors->cols, shape->type);
    
    cvGetRow(base, &submat, 0);
    cvCopy(&sh1d, &submat);
    //second row is -y0, x0, -y1, x1...
    {
	CvMat row;
	CvMat *aux = cvCreateMat(pdm.mean->rows, pdm.mean->cols, pdm.mean->type);
	cvGetRow(aux, &submat, 0);
	cvGetRow(pdm.mean, &row, 1);
	cvCopy(&row, &submat);
	cvScale(&submat, &submat, -1);

	cvGetRow(aux, &submat, 1);
	cvGetRow(pdm.mean, &row, 0);
	cvCopy(&row, &submat);

	//transpose
	CvMat *aux2 = cvCreateMat(pdm.mean->cols, pdm.mean->rows, pdm.mean->type);
	cvTranspose(aux, aux2);
	
	cvReshape(aux2, &submat, 0, 1);
	cvGetRow(base, &row, 1);
	cvCopy(&submat, &row);
	
	cvReleaseMat(&aux2);
	cvReleaseMat(&aux);	
    }

    cvGetRow(base, &submat, 2);
    for(int ai=0; ai<pdm.vectors->cols; ai+=2)
	cvmSet(&submat, 0, ai, 1);
    cvGetRow(base, &submat, 3);
    for(int ai=1; ai<pdm.vectors->cols; ai+=2)
	cvmSet(&submat, 0, ai, 1);
    

    double upd=0, upd_t1 = 10e6;
    
    unsigned int iter = 0;

    while(iter < 100){
	CvMat *residues, *weight, *hessian;
	//get the residue
	residues = get_residues(shape, pdm, params, proj);

	weight = get_robust_weight(residues);

	//calculate hessian.
	hessian = calculate_robust_hessian(weight, base);

	cvInv(hessian, ihess);
	
	//and the second term
	
	cvZero(term);
	for(int ai=0; ai<base->cols; ai++){
	    cvGetCol(base, &col, ai);
	    cvAddWeighted(term, 1, &col, 
			  cvGetReal1D(residues, ai)*cvGetReal1D(weight, ai),
			  0, term);
	}

	cvMatMul(ihess, term, update);

	//update:
	cvSub(proj, update, proj);

	cvReleaseMat(&hessian);
	cvReleaseMat(&residues);
	cvReleaseMat(&weight);

	iter++;


	upd = cvNorm(update);

	if(fabs(upd- upd_t1) < 0.05*upd_t1)
	    break;

	upd_t1 = upd;

    }

    cvReleaseMat(&base);
    cvReleaseMat(&update);
    cvReleaseMat(&term);
    cvReleaseMat(&ihess);


    return proj;

}

CvMat *robust_projection_over_shape_base(const CvMat *shape, PDModel& pdm, 
					 const CvMat *params_p0)
{


    CvMat *proj = cvCloneMat(params_p0);

    CvMat col, sh1d;
    cvReshape(shape, &sh1d, 0, 1);
    int nparams = params_p0->cols * params_p0->rows; 
   
    CvMat *ihess = cvCreateMat(nparams,nparams, shape->type);
    CvMat *term = cvCreateMat(nparams,1,shape->type);
    CvMat *update = cvCreateMat(nparams,1,shape->type);

    double upd=0, upd_t1 = 10e6;    
    unsigned int iter = 0;
    
    while(iter < 100){
	CvMat *residues, *weight, *hessian;
	//get the residue
	residues = get_residues(shape, pdm, proj, NULL);

	weight = get_robust_weight(residues);

	//calculate hessian.
	hessian = calculate_robust_hessian(weight, pdm.vectors);

	cvInv(hessian, ihess);
	
	//and the second term
	
	cvZero(term);
	for(int ai=0; ai<pdm.vectors->cols; ai++){
	    cvGetCol(pdm.vectors, &col, ai);
	    cvAddWeighted(term, 1, &col, 
			  cvGetReal1D(residues, ai)*cvGetReal1D(weight, ai),
			  0, term);
	}

	cvMatMul(ihess, term, update);

	//update:
	cvSub(proj, update, proj);

	cvReleaseMat(&hessian);
	cvReleaseMat(&residues);
	cvReleaseMat(&weight);

	iter++;

	upd = cvNorm(update);

	if(fabs(upd- upd_t1) < 0.05*upd_t1)
	    break;

	upd_t1 = upd;
    }

    cvReleaseMat(&update);
    cvReleaseMat(&term);
    cvReleaseMat(&ihess);


    return proj;
}
